{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  OchreLang.ToplevelTrans
-- Copyright   :  (c) Stephen Tetley 2014
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Translation from top level syntax.
--
--------------------------------------------------------------------------------

module OchreLang.ToplevelTrans
  ( 

    transProgram

  ) where

import OchreLang.CompilerMon
import OchreLang.Syntax
import qualified OchreLang.TopSyntax as T
import OchreLang.SyntaxBase


import Control.Applicative
import Data.List ( isPrefixOf )
import qualified Data.Map as Map
import Text.Printf



type NameEnv = Map.Map Ident VarId

find :: Ident -> NameEnv -> Maybe VarId
find = Map.lookup

extend :: Ident -> VarId -> NameEnv -> NameEnv
extend = Map.insert

extends :: [Ident] -> [VarId] -> NameEnv -> NameEnv
extends xs ys env = foldr (uncurry Map.insert) env (zip xs ys)

initial_env :: NameEnv
initial_env = Map.empty



transProgram :: T.Program -> Compiler InstrumentDict
transProgram (T.Program { T.program_modules = mod_dict }) = 
    let ms = Map.elems mod_dict 
    in fmap (buildInstrumentDict . concat) $ mapM transModule ms


transModule :: T.Module -> Compiler [Instrument]
transModule (T.Module { T.module_name = mod_name, T.module_instrs = instrs }) = 
    mapM (transInstr mod_name) instrs



transInstr :: ModuleName -> T.Instr -> Compiler Instrument
transInstr mod_name (T.Instr { T.instr_name       = name
                             , T.instr_superclass = super
                             , T.instr_methods    = ms }) = 
    do { (main_decl, tmethods)  <- splitMethods ms 
       ; main_def               <- transMainDecl main_decl
       ; methods                <- mapM transMethodDecl tmethods
       ; return $ Instrument { instr_name       = name
                             , instr_superclass = super
                             , instr_methods    = methods
                             , instr_main       = main_def
                             , instr_module     = mod_name
                             }
       }


splitMethods :: [T.MethodDecl] -> Compiler (T.MethodDecl, [T.MethodDecl])
splitMethods xs = step xs []
  where
    main_id                             = Ident "main"
    step (m@(T.MethodDecl vid _):ms) ac 
        | vid == main_id                = return (m, ac++ms)
        | otherwise                     = step ms (m:ac)
    step []                          _  = 
        throwError "ToplevelTrans.splitMethods" "Cannot find main"



transMainDecl :: T.MethodDecl -> Compiler Expr
transMainDecl (T.MethodDecl _ defn) =
     do { initNameGen
        ; MethodDefn _ e <- transMethodDefn defn
        ; return e
        }


transMethodDecl :: T.MethodDecl -> Compiler MethodDecl
transMethodDecl (T.MethodDecl vid defn) =
     do { initNameGen
        ; defn1 <- transMethodDefn defn
        ; return $ MethodDecl vid defn1
        }


transMethodDefn :: T.MethodDefn -> Compiler MethodDefn
transMethodDefn (T.MethodDefn args0 blk) = 
    do { (args,env) <- instantiateArgs initial_env args0
       ; e1         <- transDoBlock env blk
       ; return $ MethodDefn args e1
       }


instantiateArgs :: NameEnv -> [Ident] -> Compiler ([VarId], NameEnv)
instantiateArgs env []     = return $ ([],env)
instantiateArgs env (x:xs) = 
    do { vid        <- makeVarId x
       ; (vs,env1)  <- instantiateArgs (extend x vid env) xs
       ; return (vid:vs,env1)
       }

transDoBlock :: NameEnv -> T.DoBlock -> Compiler Expr
transDoBlock _   []     = throwError "ToplevelTrans.transDoBlock" "Empty 'do' block"
transDoBlock env (z:zs) = step z zs
  where
    step (T.Return e)               []      = Return <$> exprValue env e
    
    -- Ignore return if not in final position
    step (T.Return _)               (x:xs)  = step x xs


    -- Top level prim call can either be super-call, opcode call or letfun call    
    step (T.PrimCall vid es)        []      = 
        constrXCall env vid <$> mapM (exprValue env) es

    step (T.PrimCall vid es)        xs      =
        do { vs <- mapM (exprValue env) es
           ; e2 <- transDoBlock env xs
           ; let fn = constrXCall env vid
           ; return $ fn vs :>> e2
           }

    -- Bind cannot be last statement in a do-block
    step (T.DoBind _ _)             []      =
        throwError "ToplevelTrans.transDoBlock"
                   "Last statement in a 'do' block must be an expression, saw DoDind."

    step (T.DoBind patt blk)        xs      =
        do { vs         <- varBinds patt
           ; e1         <- transDoBlock env blk
           ; let env1   = extends (bindingVars patt) vs env
           ; e2         <- transDoBlock env1 xs
           ; case vs of
               [] -> return (e1 :>> e2)
               [v1] -> return $ LetSig1 v1 e1 e2
               _   -> return $ LetTuple vs e1 e2
           }

    -- Let cannot be last statement in a do-block
    step (T.DoLet {})                []     =
        throwError "ToplevelTrans.transDoBlock" 
                   "Last statement in a 'do' block must be an expression, saw DoLet."

    step (T.DoLet dc@(T.Decl v0 _)) xs      = 
        do { (vid,val)  <- transDecl env dc
           ; let env1   = extend v0 vid env
           ; expr       <- transDoBlock env1 xs
           ; return $ LetValue vid val expr
           }


    -- Lam (letfun) cannot be last statement in a do-block
    step (T.DoLam {})               []      =
        throwError "ToplevelTrans.transDoBlock" 
                   "Last statement in a 'do' block must be an expression, saw LetFun."
           
    step (T.DoLam s args d1)        xs      = 
        do { vid        <- freshLamId 
           ; vargs      <- mapM makeVarId args
           ; let env1   = extend s vid env
           ; let env2   = extends args vargs env1
           ; ea         <- transDoBlock env2 d1
           ; body       <- transDoBlock env1 xs
           ; return $ LetFun vid vargs ea body
           }

    step (T.DoIf c d1 d2)           xs      =
        do { val   <- exprValue env c
           ; e1    <- transDoBlock env d1
           ; e2    <- transDoBlock env d2
           ; let ans1 = Ifte val e1 e2
           ; if null xs then return ans1 
                        else do { rest <- transDoBlock env xs 
                                ; return $ ans1 :>> rest }
           }                                    


constrXCall :: NameEnv -> Ident -> ([Value] -> Expr)
constrXCall env ident@(Ident name) = case find ident env of
    Just lam -> FunApp lam 
    Nothing -> if "super:" `isPrefixOf` name 
               then SuperCall (Ident $ drop 6 name)
               else OpCall (Ident name)


bindingVars :: T.VarBind -> [Ident]
bindingVars (T.Bind1 v)      = [v]
bindingVars (T.BindTuple vs) = vs

varBinds :: T.VarBind -> Compiler [VarId]
varBinds = mapM makeVarId . bindingVars




transDecl :: NameEnv -> T.Decl -> Compiler (VarId,Value)
transDecl env (T.Decl vid e) = (,) <$> makeVarId vid <*> exprValue env e


exprValue :: NameEnv -> T.Expr -> Compiler Value
exprValue _   (T.Lit l)           = literalValue l


-- | A var at this position must be bound.
-- It cannot be a Csound function- or opcode name, nor an Ochre
-- method.
--
exprValue env (T.Var vid)         = case find vid env of
    Nothing -> case makeFinal (getIdent vid) of
                   Just fin -> return $ Var fin
                   Nothing -> throwError "ToplevelTrans.exprValue"  
                                          (printf "Unbound, unreconized variable: %s"
                                                  (getIdent vid))
    Just v -> return $ Var v


-- | This must be a Csound function (e.g. abs).
-- It cannot be a Csound opcode or a (pre-inlining) Ochre method
-- (Ochre methods are not pure).
--
exprValue env (T.App vid es)      =
    FunCallV vid <$> mapM (exprValue env) es

exprValue env (T.Cond ce e1 e2)   =
    CondV <$> exprValue env ce <*> exprValue env e1 <*> exprValue env e2

exprValue env (T.UnaryE op e)     = UnaryV op <$> exprValue env e

exprValue env (T.BinE op e1 e2)   =
    BinV op <$> exprValue env e1 <*> exprValue env e2

exprValue env (T.RelE op e1 e2)   =
    RelV op <$> exprValue env e1 <*> exprValue env e2

exprValue env (T.Tuple es)        = Tuple <$> mapM (exprValue env) es


literalValue :: T.Literal -> Compiler Value
literalValue (T.Bool b)         = return $ Bool b
literalValue (T.Int i)          = return $ Int i
literalValue (T.Float d)        = return $ Float d
literalValue (T.String s)       = return $ String s



makeVarId :: Ident -> Compiler VarId
makeVarId vid = case identRate vid of
    Just rate -> fresh rate
    Nothing -> throwError "ToplevelTrans.makeVarId"
                          (printf  "Unrecognized rate prefix. Cannot create a VarId for %s "
                                   (getIdent vid))



identRate :: Ident -> Maybe RateId
identRate = fn . getIdent
  where
    fn ('a':_) = Just $ A_RATE
    fn ('k':_) = Just $ K_RATE
    fn ('i':_) = Just $ I_RATE
    fn _       = Nothing


